#!/anaconda3/envs/FEALPy/bin python3.8
# -*- coding: utf-8 -*-
# File: kmeans_by_bce_embedding.py
# Author: Bryan SHEN
# E-mail: m18801919240_3@163.com
# Site: Shanghai, China
# Time: 2024/4/22 11:57
# File-Desp:

import numpy as np
import json
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
import matplotlib.pyplot as plt
from tqdm import tqdm


# def calculate_nclusters(embeddings):
#
#     # 计算不同聚类数量的轮廓系数
#     silhouette_coeffs = []
#     for k in tqdm(range(10, 30)):  # 至少需要两个聚类
#         kmeans = KMeans(n_clusters=k, random_state=0).fit(embeddings)
#         score = silhouette_score(embeddings, kmeans.labels_)
#         silhouette_coeffs.append(score)
#
#     # 绘制轮廓系数图
#     plt.figure(figsize=(10, 6))
#     plt.plot(range(31, 33), silhouette_coeffs, marker='o')
#     plt.title('Silhouette Coefficient For Each k')
#     plt.xlabel('Number of clusters')
#     plt.ylabel('Silhouette Coefficient')
#     # 保存图像到文件
#     plt.savefig('../../data/result/silhouette_coefficients.png')  # 将图片保存为PNG文件
#     plt.show()


class KMeansByBceEmbedding(object):

    def __init__(self):

        self.n_clusters = 70
        self.kmeans = KMeans(n_clusters=self.n_clusters, random_state=0)
        self.embd_file_path = '../../data/result/embeddings.jsonl'
        self.cluster_results_file = '../../data/result/cluster_result.jsonl'

    def load_embeddings(self):
        """ 加载表征处理好的 bce embedding """

        embeddings = []
        ids = []
        with open(self.embd_file_path, 'r') as f:
            for line in f:
                record = json.loads(line)
                ids.append(record['id'])
                embeddings.append(record['embedding'])

        return np.array(embeddings), ids

    def calculate_nclusters(self, embeddings):
        """ 基于不同的embedding计算不同聚类数量的轮廓系数 """

        silhouette_coeffs = []
        for k in tqdm(range(10, 30)):  # 至少需要两个聚类
            kmeans = KMeans(n_clusters=k, random_state=0).fit(embeddings)
            score = silhouette_score(embeddings, kmeans.labels_)
            silhouette_coeffs.append(score)

        # 绘制轮廓系数图
        plt.figure(figsize=(10, 6))
        plt.plot(range(31, 33), silhouette_coeffs, marker='o')
        plt.title('Silhouette Coefficient For Each k')
        plt.xlabel('Number of clusters')
        plt.ylabel('Silhouette Coefficient')
        # 保存图像到文件
        plt.savefig('../../data/result/silhouette_coefficients.png')  # 将图片保存为PNG文件
        plt.show()

    def run_cluster(self):

        embeddings, sentence_ids = k.load_embeddings()

        self.kmeans.fit(embeddings)
        cluster_labels = self.kmeans.labels_

        # Save clustering results to a JSONL file
        with open(self.cluster_results_file, 'w') as file:
            for id, label in zip(sentence_ids, cluster_labels):
                json_record = json.dumps({"id": int(id), "cluster_label": int(label)})
                file.write(json_record + '\n')


if __name__ == '__main__':

    embeddings_file = '../../data/result/embeddings.jsonl'

    k = KMeansByBceEmbedding()
    # k.calculate_nclusters()
    # embeddings, sentence_ids = k.load_embeddings()
    k.run_cluster()

    # Load embeddings from file
    # embeddings, sentence_ids = load_embeddings(embeddings_file)



